import os, csv, shutil

import torch
from torchvision.datasets.vision import VisionDataset

from PIL import Image


class CassavaDataset(VisionDataset):
    def __init__(self, root, img_list_file, istrain=True, transforms=None, transform=None, target_transform=None):
        super(CassavaDataset, self).__init__(root, transforms, transform, target_transform)
        # self.classes = os.listdir(path)
        # self.path = [f"{path}/{className}" for className in self.classes]
        # self.file_list = [glob.glob(f"{x}/*") for x in self.path]
        # self.transform = transform
        self.file_list = []
        if istrain:
            with open(os.path.join(root, img_list_file))as f:
                f_csv = csv.reader(f)
                for row in f_csv:
                    row[0] = os.path.join(root, "train_images", row[0])
                    times = 1
                    if row[1] == '0':
                        times = 11
                    elif row[1] == '1':
                        times = 6
                    elif row[1] == '2':
                        times = 5
                    elif row[1] == '4':
                        times = 5
                    for t in range(times):
                        self.file_list.append(row)
        else:
            with open(os.path.join(root, img_list_file))as f:
                f_csv = csv.reader(f)
                for row in f_csv:
                    row[0] = os.path.join(root, "train_images", row[0])
                    self.file_list.append(row)
        # for i, className in enumerate(self.classes):
        #     for fileName in self.file_list[i]:
        #         files.append([i, className, fileName])

        # files = None

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):
        fileName = self.file_list[idx][0]
        label = int(self.file_list[idx][1])
        img = Image.open(fileName)
        if self.transform is not None:
            img = self.transform(img)

        return img, label, fileName


if __name__ == '__main__':
    root = "/home/handewei/data/cassava-leaf/"

    trainSet = CassavaDataset(root)
    for i in range(trainSet.__len__()):
        img, label, fileName = trainSet[i]
        img = Image.Image.resize(img,(256, 256))
        save_path = os.path.join(root, "ana", str(label), (str(i) + ".jpg"))
        Image.Image.save(img,save_path)
        # shutil.copy(fileName, save_path)
        if(i%100==0):
            print(i)


